from Crypto.Cipher import AES
import base64
import hashlib

class AESCipher:
    def __init__(self, key):
        self.block_size = AES.block_size  # 16
        self.key = hashlib.sha256(key.encode('utf-8')).digest()  # 32 字节密钥

    def pad(self, data):
        padding_length = self.block_size - len(data.encode('utf-8')) % self.block_size
        padding = chr(padding_length) * padding_length
        return data + padding

    def unpad(self, data):
        padding_length = ord(data[-1])
        return data[:-padding_length]

    def encrypt(self, raw):
        raw_padded = self.pad(raw)
        cipher = AES.new(self.key, AES.MODE_ECB)
        encrypted = cipher.encrypt(raw_padded.encode('utf-8'))
        return base64.urlsafe_b64encode(encrypted).decode('utf-8')  # 改为 URL 安全的编码

    def decrypt(self, enc):
        cipher = AES.new(self.key, AES.MODE_ECB)
        decrypted = cipher.decrypt(base64.urlsafe_b64decode(enc)).decode('utf-8')
        return self.unpad(decrypted)
